# -*- coding: utf-8 -*-
'''
Created on 17.10.2019

@author: yu03
'''
import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from Video_Unified import folder_path, np_result_names, hor_index, ver_index, hor_lines, ver_lines
import sys
from FFT_Interpolation import FFT_interpolation_nonlinearity_compare
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

plt_show_mode = 'show' ### plt.show()
# plt_show_mode = 'save' ### plt.save()

i = 0
fs_cam = 50
''' 
    for all groups
''' 
for np_result in np_result_names:
    ''' 
        reading .npy file
    ''' 
    file_name = np_result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
    i += 1
    f = open(np_result, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    ''' 
        put 4-D results in "lines"
    ''' 
    lines = []
    time_sequence = []
    while f.tell() < f_size:
        line = np.load(f, allow_pickle=True)
        lines.append(line)
    print('%i: %s, %eB'%(i, file_name, f_size))
#     lines = np.array(lines)[::,::,300:2200]
    lines = np.array(lines)[::,::,300:800]
    print(np.shape(lines), 'Channle, lines, frames:')
    frame_num = np.shape(lines)[2]
    line_num = np.shape(lines)[1]
    
    hor_angle_set, ver_angle_set, hor_length_set, ver_length_set = lines
    ''' 
        making average
    ''' 
    averaged_results = []
    
    hor_angle_avr = np.mean(hor_angle_set, axis=0)
    ver_angle_avr = np.mean(ver_angle_set, axis=0)
    hor_length_avr = np.mean(hor_length_set, axis=0)
    ver_length_avr = np.mean(ver_length_set, axis=0)
    
    averaged_results.append(hor_angle_avr)
    averaged_results.append(ver_angle_avr)
    averaged_results.append(hor_length_avr)
    averaged_results.append(ver_length_avr)
        
    ''' 
        polyfit degree
    '''     
    length_dof = 10
    angle_dof = 6

    ''' 
        fitting averaged results
    '''
    averaged_nonlinearity = []
    for data in averaged_results[0:2]:
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        nonlinearity = data - data_fitline
        averaged_nonlinearity.append(nonlinearity)
    for data in averaged_results[2:4]:
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        nonlinearity = data - data_fitline
        averaged_nonlinearity.append(nonlinearity)

    print(np.shape(averaged_nonlinearity))
    [hor_angle_avr_nonlinearity, ver_angle_avr_nonlinearity, hor_length_avr_nonlinearity, ver_length_avr_nonlinearity] = averaged_nonlinearity
    print(np.shape(hor_angle_avr_nonlinearity))

    ''' 
        fitting linebyline results
    '''
    hor_angle_nonlinearity_set = []
    ver_angle_nonlinearity_set = []
    hor_length_nonlinearity_set = []
    ver_length_nonlinearity_set = []    
      
    for i in range(line_num):
#         for i in [50,51,52]:
        data = hor_angle_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        hor_angle_nonlinearity = data - data_fitline
           
           
        data = hor_length_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        hor_length_nonlinearity = data - data_fitline
           
        data = ver_angle_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, angle_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        ver_angle_nonlinearity = data - data_fitline
           
        data = ver_length_set[i]
        data_fit_params = np.polyfit(np.arange(frame_num), data, length_dof)
        data_fit_poly = np.poly1d(data_fit_params)
        data_fitline = data_fit_poly(np.arange(frame_num))
        ver_length_nonlinearity = data - data_fitline
           
        hor_angle_nonlinearity_set.append(hor_angle_nonlinearity)
        hor_length_nonlinearity_set.append(hor_length_nonlinearity)
        ver_angle_nonlinearity_set.append(ver_angle_nonlinearity)
        ver_length_nonlinearity_set.append(ver_length_nonlinearity)
           
           
    hor_angle_nonlinearity_set = np.array(hor_angle_nonlinearity_set)
    hor_length_nonlinearity_set = np.array(hor_length_nonlinearity_set)
    ver_angle_nonlinearity_set = np.array(ver_angle_nonlinearity_set)
    ver_length_nonlinearity_set = np.array(ver_length_nonlinearity_set)
       
       
    hor_angle_nonlinearity_max = np.amax(hor_angle_nonlinearity_set, axis=1)
    hor_angle_nonlinearity_min = np.amin(hor_angle_nonlinearity_set, axis=1)
    hor_length_nonlinearity_max = np.amax(hor_length_nonlinearity_set, axis=1)
    hor_length_nonlinearity_min = np.amin(hor_length_nonlinearity_set, axis=1)
       
    ver_angle_nonlinearity_max = np.amax(ver_angle_nonlinearity_set, axis=1)
    ver_angle_nonlinearity_min = np.amin(ver_angle_nonlinearity_set, axis=1)
    ver_length_nonlinearity_max = np.amax(ver_length_nonlinearity_set, axis=1)
    ver_length_nonlinearity_min = np.amin(ver_length_nonlinearity_set, axis=1)
#     ''' 
#         Nonlinearity FFT
#     ''' 
#     fig = plt.figure('Nonlinearity FFT')
#     plt.gcf().set_size_inches(18,9)
#     ax1 = fig.add_subplot(2, 2, 1, projection='3d')
#     ax2 = fig.add_subplot(2, 2, 2, projection='3d')
#     ax3 = fig.add_subplot(2, 2, 3, projection='3d')
#     ax4 = fig.add_subplot(2, 2, 4, projection='3d')
#     plt.gca().patch.set_facecolor('white')
#     
#     hor_non_fft_set = np.empty(shape=[101, 150])
#     
#     for i in np.arange(101):
#         hor_non_f, hor_non_phi, freqline, fft = FFT_interpolation_nonlinearity_compare(hor_angle_nonlinearity_set[int(i), 200:1000], 1/fs_cam)
#         hor_non_fft_set[i] = np.array(fft[10:160])
#     hor_non_fft_set = np.array(hor_non_fft_set)
#     X_fft = np.arange(101)
#     Y_fft = freqline[10:160]
#     X_fft, Y_fft = np.meshgrid(Y_fft, X_fft)
#     surf1 = ax1.plot_surface(X_fft, Y_fft, hor_non_fft_set, cmap='jet', rstride=1, cstride=1,
#                            linewidth=0, antialiased=False)
#     ax1.yaxis.set_major_formatter(FormatStrFormatter('%i'))
#        
#     if plt_show_mode == 'save':
# #         plt.savefig(folder_path + '\\' + 'linebyline_Raw_Data.png', dpi=300)
#         plt.close()
#     elif plt_show_mode == 'show':
#         plt.show()
#     else:
#         sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)   
#     
    ''' 
        Plotting
    '''   
    y = [np.arange(frame_num)]*line_num
    y = np.array(y)
    
    ''' 
        Raw Data
    ''' 
    fig = plt.figure('linebyline Raw Data')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = fig.add_subplot(2, 2, 1, projection='3d')
    ax2 = fig.add_subplot(2, 2, 2, projection='3d')
    ax3 = fig.add_subplot(2, 2, 3, projection='3d')
    ax4 = fig.add_subplot(2, 2, 4, projection='3d')
    plt.gca().patch.set_facecolor('white')
           
       
    x = [hor_lines]*frame_num
    x = np.array(x).T
    for i in range(line_num):
    #         if i%10 == 0:
        ax1.plot(x[i,0:-1],y[i,0:-1],hor_angle_set[i,0:-1])
    for i in range(line_num):
        ax2.plot(x[i,:],y[i,:],hor_length_set[i,:])
       
    x = [ver_lines]*frame_num
    x = np.array(x).T
    for i in range(line_num):
        ax3.plot(x[i,:],y[i,:],ver_angle_set[i,:])
    for i in range(line_num):
        ax4.plot(x[i,:],y[i,:],ver_length_set[i,:])
       
       
    ax1.title.set_text('Horizontal Tilt')
    ax2.title.set_text('Horizontal Length')
    ax3.title.set_text('Vertical Tilt')
    ax4.title.set_text('Vertical Length')
       
    ax1.set_xlabel('Hor. Line Num.')
    ax1.set_ylabel('Frame Num.')
    ax1.set_zlabel('Hor. Tilt / urad')
       
    ax2.set_xlabel('Hor. Line Num.')
    ax2.set_ylabel('Frame Num.')
    ax2.set_zlabel('Hor. Length / nm')
       
    ax3.set_xlabel('Ver. Line Num.')
    ax3.set_ylabel('Frame Num.')
    ax3.set_zlabel('Ver. Tilt / urad')
       
    ax4.set_xlabel('Ver. Line Num.')
    ax4.set_ylabel('Frame Num.')
    ax4.set_zlabel('Ver. Length / nm')
       
    ax1.view_init(0, 0)
    ax2.view_init(0, 0)
    ax3.view_init(0, 0)
    ax4.view_init(0, 0)
       
    # figManager = plt.get_current_fig_manager()
    # figManager.window.showMaximized()
    plt.tight_layout()
    
    if plt_show_mode == 'save':
        plt.savefig(folder_path + '\\' + 'linebyline_Raw_Data.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        pass
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
    
    
    ''' 
        linebyline Nonlinearity Zoom-In
    ''' 
    fig = plt.figure('linebyline Nonlinearity Zoom-In')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = fig.add_subplot(2, 2, 1, projection='3d')
    ax2 = fig.add_subplot(2, 2, 2, projection='3d')
    ax3 = fig.add_subplot(2, 2, 3, projection='3d')
    ax4 = fig.add_subplot(2, 2, 4, projection='3d')
    plt.gca().patch.set_facecolor('white')
    zoom_start, zoom_end = 0, -1
    zoom_start, zoom_end = 200, 1000
       
    x = [hor_lines]*frame_num
    x = np.array(x).T
    line_zoom = np.linspace(0,100,num=16,dtype=int)
    #     for i in range(line_num):
    for i in line_zoom:
        ax1.plot(x[int(i),zoom_start:zoom_end],y[int(i),zoom_start:zoom_end], hor_angle_nonlinearity_set[int(i), zoom_start:zoom_end])
    #     for i in range(line_num):
    for i in line_zoom:
        ax2.plot(x[int(i),zoom_start:zoom_end],y[int(i),zoom_start:zoom_end],hor_length_nonlinearity_set[int(i), zoom_start:zoom_end])
       
    x = [ver_lines]*frame_num
    x = np.array(x).T
    #     for i in range(line_num):
    for i in line_zoom:
        ax3.plot(x[int(i),zoom_start:zoom_end],y[int(i),zoom_start:zoom_end],ver_angle_nonlinearity_set[int(i), zoom_start:zoom_end])
    #         ax3.plot(x[i,:],y[i,:],ver_angle_set[i,:])
    #     for i in range(line_num):
    for i in line_zoom:
        ax4.plot(x[int(i),zoom_start:zoom_end],y[int(i),zoom_start:zoom_end],ver_length_nonlinearity_set[int(i), zoom_start:zoom_end])
    #             ax4.plot(x[i,:],y[i,:],ver_length_set[i,:])
       
    ax1.title.set_text('Horizontal Tilt')
    ax2.title.set_text('Horizontal Length')
    ax3.title.set_text('Vertical Tilt')
    ax4.title.set_text('Vertical Length')
       
    ax1.set_xlabel('Hor. Line Num.')
    ax1.set_ylabel('Frame Num.')
    ax1.set_zlabel('Hor. Tilt / urad')
       
    ax2.set_xlabel('Hor. Line Num.')
    ax2.set_ylabel('Frame Num.')
    ax2.set_zlabel('Hor. Length / nm')
       
    ax3.set_xlabel('Ver. Line Num.')
    ax3.set_ylabel('Frame Num.')
    ax3.set_zlabel('Ver. Tilt / urad')
       
    ax4.set_xlabel('Ver. Line Num.')
    ax4.set_ylabel('Frame Num.')
    ax4.set_zlabel('Ver. Length / nm')
       
    ax1.view_init(60, 60)
    ax2.view_init(60, 60)
    ax3.view_init(80, 0)
    ax4.view_init(80, 0)
       
    # figManager = plt.get_current_fig_manager()
    # figManager.window.showMaximized()
    plt.tight_layout()
    #     plt.show()
    if plt_show_mode == 'save':
        plt.savefig(folder_path + '\\' + 'linebyline_Zoom_In.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        pass
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
        
    
    
    ''' 
        linebyline Nonlinearity Amplitude
    ''' 
    fig = plt.figure('linebyline Nonlinearity Amplitude')
    plt.gcf().set_size_inches(18,9)
    plt.subplot(2,2,1)
    plt.plot(np.arange(line_num), hor_angle_nonlinearity_max, 'b')
    plt.plot(np.arange(line_num), hor_angle_nonlinearity_min, 'b')
    x = np.concatenate((np.arange(line_num),np.arange(line_num)[::-1]))
    y = np.concatenate((hor_angle_nonlinearity_max, hor_angle_nonlinearity_min[::-1]))
    plt.fill(x, y,facecolor='g',alpha=0.5)
    plt.title('Hor. Tilt')
    plt.xlabel('Line Num.')
    plt.ylabel('urad')
    plt.grid(which='both', axis='both')
       
    plt.subplot(2,2,2)
    plt.plot(np.arange(line_num), hor_length_nonlinearity_max, 'b')
    plt.plot(np.arange(line_num), hor_length_nonlinearity_min, 'b')
    x = np.concatenate((np.arange(line_num),np.arange(line_num)[::-1]))
    y = np.concatenate((hor_length_nonlinearity_max, hor_length_nonlinearity_min[::-1]))
    plt.fill(x, y,facecolor='g',alpha=0.5)
    plt.title('Hor. Length')
    plt.xlabel('Line Num.')
    plt.ylabel('nm')
    plt.grid(which='both', axis='both')
       
    plt.subplot(2,2,3)
    plt.plot(np.arange(line_num), ver_angle_nonlinearity_max, 'b')
    plt.plot(np.arange(line_num), ver_angle_nonlinearity_min, 'b')
    x = np.concatenate((np.arange(line_num),np.arange(line_num)[::-1]))
    y = np.concatenate((ver_angle_nonlinearity_max, ver_angle_nonlinearity_min[::-1]))
    plt.fill(x, y,facecolor='g',alpha=0.5)
    plt.title('Ver. Tilt')
    plt.xlabel('Line Num.')
    plt.ylabel('urad')
    plt.grid(which='both', axis='both')
       
    plt.subplot(2,2,4)
    plt.plot(np.arange(line_num), ver_length_nonlinearity_max, 'b')
    plt.plot(np.arange(line_num), ver_length_nonlinearity_min, 'b')
    x = np.concatenate((np.arange(line_num),np.arange(line_num)[::-1]))
    y = np.concatenate((ver_length_nonlinearity_max, ver_length_nonlinearity_min[::-1]))
    plt.fill(x, y,facecolor='g',alpha=0.5)
    plt.title('Ver. Length')
    plt.xlabel('Line Num.')
    plt.ylabel('nm')
    plt.grid(which='both', axis='both')
       
    # figManager = plt.get_current_fig_manager()
    # figManager.window.showMaximized()
    plt.tight_layout()
    
    if plt_show_mode == 'save':
        plt.savefig(folder_path + '\\' + 'linebyline_Nonlinearity_Amplitude.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        pass
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
    
    ''' 
        T_based Nonlinearity Averaged
    ''' 
    fig = plt.figure('T_based Nonlinearity Averaged')
    plt.gcf().set_size_inches(18,9)
    plt.subplot(2,2,1)
    plt.plot(np.arange(frame_num), hor_angle_nonlinearity_set[51], color='blue', label='single-line')
    plt.plot(np.arange(frame_num), hor_angle_avr_nonlinearity, color='red', label='averaged')
    plt.title('Hor. Tilt')
    plt.xlabel('Frame Num.')
    plt.ylabel('urad')
    plt.grid(which='both', axis='both')
    plt.legend()
    
    plt.subplot(2,2,2)
    plt.plot(np.arange(frame_num), ver_angle_nonlinearity_set[51], color='blue', label='single-line')
    plt.plot(np.arange(frame_num), ver_angle_avr_nonlinearity, color='red', label='averaged')
    plt.title('Ver. Tilt')
    plt.xlabel('Frame Num.')
    plt.ylabel('urad')
    plt.grid(which='both', axis='both')
    plt.legend()
    
    plt.subplot(2,2,3)
    plt.plot(np.arange(frame_num), hor_length_nonlinearity_set[51], color='blue', label='single-line')
    plt.plot(np.arange(frame_num), hor_length_avr_nonlinearity, color='red', label='averaged')
    plt.title('Hor. Length')
    plt.xlabel('Frame Num.')
    plt.ylabel('nm')
    plt.grid(which='both', axis='both')
    plt.legend()
    
    plt.subplot(2,2,4)
    plt.plot(np.arange(frame_num), ver_length_nonlinearity_set[51], color='blue', label='single-line')
    plt.plot(np.arange(frame_num), ver_length_avr_nonlinearity, color='red', label='averaged')
    plt.title('Ver. Length')
    plt.xlabel('Frame Num.')
    plt.ylabel('nm')
    plt.grid(which='both', axis='both')
    plt.legend()
    
    
    # figManager = plt.get_current_fig_manager()
    # figManager.window.showMaximized()
    plt.tight_layout()
    
    if plt_show_mode == 'save':
        plt.savefig(folder_path + '\\' + 'Nonlinearity(T_based)_averaging.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        plt.show()
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)


